package com.sz.admin.ai.factory.embeddingFactory.entity;


import com.sz.admin.ai.factory.embeddingFactory.Handler.impl.OllamaEmbeddingHandlerImpl;
import lombok.Getter;

import java.util.List;
import java.util.Objects;

/**
 * @描述: 嵌入模型平台枚举
 * @Date: 2025/1/10 09:33
 * @Author: TangYuan
 * @Date: 2025/1/10 09:33
 */
@Getter
public enum EmbeddingPlatformEnum {
    
    //? 1. ollama平台 --------------------------------------------------------------------------------------------------

    OLLAMA(
            "ollama",
            List.of(
                    new EmbeddingModelEntity("nomic-embed-text:latest", 768),
                    new EmbeddingModelEntity("bge-m3:latest", 1024)
            ),
            OllamaEmbeddingHandlerImpl.class
    ),
    ALIYUN(
            "aliyun",
            List.of(
                    new EmbeddingModelEntity("text-embedding-v3", 1024),
                    new EmbeddingModelEntity("text-embedding-v2", 1536),
                    new EmbeddingModelEntity("text-embedding-v1", 1536)
            ),
            null
    ),
    
    //? 1. ollama平台 --------------------------------------------------------------------------------------------------
    ;
    
    // 平台名称
    // 模型名称
    // 维度
    
    // 平台名称
    private final String platformName;
    
    // 模型名称
    private final List<EmbeddingModelEntity> modelList;
    
    // 处理类
    private final Class<?> handlerClass;
    
    
    /**
     * 构造函数
     * @param platformName
     * @param modelList
     * @param handlerClass
     * @return {@link null}
     */
    EmbeddingPlatformEnum(String platformName, List<EmbeddingModelEntity> modelList, Class<?> handlerClass) {
        this.platformName = platformName;
        this.modelList = modelList;
        this.handlerClass = handlerClass;
    }
    
    /**
     * 根据平台名称获取嵌入模型平台枚举
     * @param platformName 平台名称
     * @return {@link EmbeddingPlatformEnum}
     */
    public static EmbeddingPlatformEnum getEnumByPlatformName(String platformName) {
        for (EmbeddingPlatformEnum embeddingPlatformEnum : EmbeddingPlatformEnum.values()) {
            if (Objects.equals(embeddingPlatformEnum.platformName, platformName)) {
                return embeddingPlatformEnum;
            }
        }
        return null;
    }
}
